/*
 * ScreenRecord - A screen recording tool for OpenHarmony devices
 * Copyright (C) 2025 Frank Zhang <frankzhang2010@foxmail.com>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 */

#include "hilog/log.h"
#include <multimedia/player_framework/native_avscreen_capture.h>
#include <multimedia/player_framework/native_avscreen_capture_base.h>
#include <multimedia/player_framework/native_avcodec_videoencoder.h>
#include <multimedia/player_framework/native_avformat.h>
#include <multimedia/player_framework/native_avmemory.h>
#include <unistd.h>
#include <csignal>
#include <string>
#include <getopt.h>

#define LOG_TAG "ScreenRecord"


void OnError(OH_AVScreenCapture *capture, int32_t errorCode, void *userData) {
    fprintf(stderr, "screenrecord OnError: %d\n", errorCode);
}

void OnStateChange(struct OH_AVScreenCapture *capture, OH_AVScreenCaptureStateCode stateCode, void *userData) {
    if (stateCode == OH_SCREEN_CAPTURE_STATE_STARTED) {
        OH_LOG_INFO(LOG_APP, "Screen capture started");
    }
}

struct OH_AVScreenCapture *capture;
OH_AVCodec *h264Encoder = nullptr;

// 全局变量用于控制录制
static bool g_isRecording = true;
static FILE* g_outputFile = nullptr;
static bool g_useStdout = false;

// 参数默认值
struct ScreenRecordOptions {
    std::string fileName = "/data/local/tmp/video.mp4";
    int width = 720;
    int height = 1280;
    int frameRate = 30;
    int timeLimit = 0;  // 0表示无限制
};

// 信号处理函数
void SignalHandler(int signo) {
    if (signo == SIGINT || signo == SIGTERM) {
        g_isRecording = false;
        OH_LOG_INFO(LOG_APP, "Received signal %d, stopping recording...", signo);
    }
}

void PrintUsage() {
    fprintf(stderr,
            "Usage: screenrecord [options]\n"
            "Options:\n"
            "--file FILE         Output file path (use '-' for stdout)\n"
            "--width WIDTH       Set frame width (64-1920, default: 720)\n"
            "--height HEIGHT     Set frame height (64-1920, default: 1280)\n"
            "--framerate FPS     Set frames per second (1-60, default: 30)\n"
            "--time-limit SEC    Set maximum recording time in seconds\n"
            "                    Range: 0-3600, 0 means unlimited (default: 0)\n"
            "\nExample:\n"
            "  screenrecord --width 1280 --height 720 --framerate 30 "
            "--time-limit 60 --file /data/video.mp4\n");
}

// 修改OnEncoderFrameAvailable函数
void OnEncoderFrameAvailable(OH_AVCodec *codec, uint32_t index, OH_AVMemory *data,
                            OH_AVCodecBufferAttr *attr, void *userData) {
    if (attr && attr->size > 0 && g_isRecording) {
        uint8_t *buffer = static_cast<uint8_t*>(OH_AVMemory_GetAddr(data));
        if (g_useStdout) {
            fwrite(buffer, 1, attr->size, stdout);
            fflush(stdout);
        } else if (g_outputFile) {
            fwrite(buffer, 1, attr->size, g_outputFile);
            fflush(g_outputFile);
        }
        OH_LOG_INFO(LOG_APP, "Encoded frame Size: %{public}d PTS: %{public}lld",
                    attr->size, attr->pts);
    }
    OH_VideoEncoder_FreeOutputData(codec, index);
}

// 添加参数校验函数
bool ValidateOptions(const ScreenRecordOptions& options) {
    // 检查文件名
    if (options.fileName.empty() && !g_useStdout) {
        fprintf(stderr, "Output file name cannot be empty\n");
        return false;
    }

    // 检查分辨率范围
    if (options.width < 64) {
        fprintf(stderr, "Invalid width: %d (must be >= 64)\n",
                options.width);
        return false;
    }
    if (options.height < 64) {
        fprintf(stderr, "Invalid height: %d (must be >= 64)\n",
                options.height);
        return false;
    }

    // 检查帧率范围
    if (options.frameRate < 10 || options.frameRate > 120) {
        fprintf(stderr, "Invalid frame rate: %d (must be between 10 and 120)\n",
                options.frameRate);
        return false;
    }

    // 检查时间限制
    if (options.timeLimit < 0) {
        fprintf(stderr, "Invalid time limit: %d (must be >= 0)\n",
                options.timeLimit);
        return false;
    }

    return true;
}

int main(int argc, char *argv[]) {
    ScreenRecordOptions options;

    // 解析命令行参数
    static struct option long_options[] = {
        {"file", required_argument, nullptr, 'f'},
        {"width", required_argument, nullptr, 'w'},
        {"height", required_argument, nullptr, 'h'},
        {"framerate", required_argument, nullptr, 'r'},
        {"time-limit", required_argument, nullptr, 't'},
        {nullptr, 0, nullptr, 0}
    };

    int opt;
    while ((opt = getopt_long(argc, argv, "", long_options, nullptr)) != -1) {
        switch (opt) {
            case 'f':
                options.fileName = optarg;
                if (options.fileName == "-") {
                    g_useStdout = true;
                }
                break;
            case 'w':
                options.width = atoi(optarg);
                break;
            case 'h':
                options.height = atoi(optarg);
                break;
            case 'r':
                options.frameRate = atoi(optarg);
                break;
            case 't':
                options.timeLimit = atoi(optarg);
                break;
            default:
                PrintUsage();
                return -1;
        }
    }

    // 添加参数校验
    if (!ValidateOptions(options)) {
        PrintUsage();
        return -1;
    }

    // 打印当前配置
    OH_LOG_INFO(LOG_APP, "Configuration:");
    OH_LOG_INFO(LOG_APP, "  Output: %{public}s",
                g_useStdout ? "stdout" : options.fileName.c_str());
    OH_LOG_INFO(LOG_APP, "  Video size: %{public}dx%{public}d",
                options.width, options.height);
    OH_LOG_INFO(LOG_APP, "  Frame rate: %{public}d fps", options.frameRate);
    OH_LOG_INFO(LOG_APP, "  Time limit: %{public}d seconds", options.timeLimit);

    // 设置信号处理
    signal(SIGINT, SignalHandler);
    signal(SIGTERM, SignalHandler);

    OH_LOG_INFO(LOG_APP, "screenrecord started!");

    // 打开输出文件
    if (!g_useStdout) {
        g_outputFile = fopen(options.fileName.c_str(), "wb");
        if (!g_outputFile) {
            fprintf(stderr, "Failed to open output file\n");
            return -1;
        }
    }

    capture = OH_AVScreenCapture_Create();
    OH_AVScreenCapture_SetErrorCallback(capture, OnError, nullptr);
    OH_AVScreenCapture_SetStateCallback(capture, OnStateChange, nullptr);

    OH_VideoCaptureInfo videocapinfo = {
        .videoFrameWidth = options.width,
        .videoFrameHeight = options.height,
        .videoSource = OH_VIDEO_SOURCE_SURFACE_RGBA
    };
    OH_VideoInfo videoinfo = {.videoCapInfo = videocapinfo};
    OH_AVScreenCaptureConfig config = {
        .captureMode = OH_CAPTURE_HOME_SCREEN,
        .dataType = OH_ORIGINAL_STREAM,
        .videoInfo = videoinfo
    };
    OH_AVScreenCapture_Init(capture, config);

    h264Encoder = OH_VideoEncoder_CreateByMime(OH_AVCODEC_MIMETYPE_VIDEO_AVC);
    if (h264Encoder == nullptr) {
        fprintf(stderr, "Failed to create H.264 encoder\n");
        return -1;
    }

    OH_AVFormat *format = OH_AVFormat_Create();
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_WIDTH, options.width);
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_HEIGHT, options.height);
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_PIXEL_FORMAT, AV_PIXEL_FORMAT_RGBA);
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_FRAME_RATE, options.frameRate);
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_I_FRAME_INTERVAL, 1);
    OH_AVFormat_SetIntValue(format, OH_MD_KEY_PROFILE, AVC_PROFILE_MAIN);

    OH_AVCodecAsyncCallback cb = {
        .onError = [](OH_AVCodec *codec, int32_t errorCode, void *userData) {
            fprintf(stderr, "Encoder error: %d\n", errorCode);
        },
        .onNeedOutputData = OnEncoderFrameAvailable,
    };
    OH_VideoEncoder_SetCallback(h264Encoder, cb, nullptr);

    int ret = OH_VideoEncoder_Configure(h264Encoder, format);
    OH_AVFormat_Destroy(format);
    if (ret != AV_ERR_OK) {
        fprintf(stderr, "Failed to configure encoder\n");
        return -1;
    }

    OHNativeWindow *window;
    ret = OH_VideoEncoder_GetSurface(h264Encoder, &window);
    if (ret != AV_ERR_OK) {
        fprintf(stderr, "Failed to get surface\n");
        return -1;
    }

    ret = OH_VideoEncoder_Start(h264Encoder);
    if (ret != AV_ERR_OK) {
        fprintf(stderr, "Failed to start encoder\n");
        return -1;
    }

    ret = OH_AVScreenCapture_StartScreenCaptureWithSurface(capture, window);
    if (ret != AV_SCREEN_CAPTURE_ERR_OK) {
        fprintf(stderr, "Failed to start screen capture\n");
        return -1;
    }

    // 修改主循环，支持时间限制
    time_t startTime = time(nullptr);
    while (g_isRecording) {
        if (options.timeLimit > 0 &&
            (time(nullptr) - startTime) >= options.timeLimit) {
            OH_LOG_INFO(LOG_APP, "Time limit reached");
            break;
        }
        sleep(1);
    }

    // 清理资源
    OH_LOG_INFO(LOG_APP, "Cleaning up...");

    if (capture) {
        OH_AVScreenCapture_StopScreenCapture(capture);
        OH_AVScreenCapture_Release(capture);
    }

    if (h264Encoder) {
        OH_VideoEncoder_Stop(h264Encoder);
        OH_VideoEncoder_Destroy(h264Encoder);
    }

    if (g_outputFile) {
        fclose(g_outputFile);
    }

    OH_LOG_INFO(LOG_APP, "Recording finished");
    return 0;
}
